# coding=utf-8
# v0.1.0 host文件解析，新增host line，删除域名

import platform
import re
import argparse

host_line_pattern = re.compile(r'^(\S+)\s+([^#]+)(#.+)?$')

ipv6_pattern = re.compile(
    r'^((([0-9A-Fa-f]{1,4}:){7}([0-9A-Fa-f]{1,4}|:))'
    r'|(([0-9A-Fa-f]{1,4}:){6}(:[0-9A-Fa-f]{1,4}'
    r'|((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))'
    r'|(([0-9A-Fa-f]{1,4}:){5}(((:[0-9A-Fa-f]{1,4}){1,2})'
    r'|:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))'
    r'|(([0-9A-Fa-f]{1,4}:){4}(((:[0-9A-Fa-f]{1,4}){1,3})'
    r'|((:[0-9A-Fa-f]{1,4})?:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))'
    r'|(([0-9A-Fa-f]{1,4}:){3}(((:[0-9A-Fa-f]{1,4}){1,4})'
    r'|((:[0-9A-Fa-f]{1,4}){0,2}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))'
    r'|(([0-9A-Fa-f]{1,4}:){2}(((:[0-9A-Fa-f]{1,4}){1,5})'
    r'|((:[0-9A-Fa-f]{1,4}){0,3}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))'
    r'|(([0-9A-Fa-f]{1,4}:)(((:[0-9A-Fa-f]{1,4}){1,6})'
    r'|((:[0-9A-Fa-f]{1,4}){0,4}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))'
    r'|(:(((:[0-9A-Fa-f]{1,4}){1,7})'
    r'|((:[0-9A-Fa-f]{1,4}){0,5}:((25[0-5]|2[0-4]\d|1\d\d'
    r'|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:)))(%.+)?$')

ipv4_pattern = re.compile(
    r'^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$')

blank_pattern = re.compile(r'\s+')


def is_comment(text):
    return text.strip().startswith('#')


class HostLine(object):

    def __init__(self, line):
        self.__text = ''
        self.__domains = set()
        self.__ip = ''
        self.__comment = ''
        self.__is_comment = True
        self.__modified = False
        self.parse(line)

    @property
    def ip(self):
        return self.__ip

    @property
    def is_comment(self):
        return self.__is_comment

    @property
    def domains(self):
        return set(self.__domains)

    def parse(self, line):

        self.__text = line
        self.__is_comment = is_comment(line)
        match = host_line_pattern.match(line)

        if not self.__is_comment:
            if match:
                ip = match.group(1)
                domains = match.group(2)
                comment = match.group(3) or ''
                if domains and ipv6_pattern.match(ip) or ipv4_pattern.match(ip):
                    self.__ip = ip
                    self.__domains = set(blank_pattern.split(domains))
                    self.__comment = comment
                    return
            else:
                self.__text = '# ' + line

        self.__is_comment = True

    def contains_domain(self, domain):
        return domain in self.__domains

    def _modified(self):
        self.__modified = True
        if not self.__domains or not self.__ip:
            if not is_comment(self.__text):
                self.__text = '# ' + self.__text
                self.__is_comment = True

    def remove_domain(self, domain):
        if self.contains_domain(domain):
            self.__domains.remove(domain)
            self._modified()

    def add_domain(self, domain):
        if not self.contains_domain(domain):
            self.__domains.add(domain)
            self._modified()

    def __str__(self):
        if self.__is_comment or not self.__modified:
            return self.__text
        return '%s\t%s%s' % (self.__ip, " ".join(self.__domains), self.__comment)


class Hosts(object):

    def __init__(self):
        self.__hosts_lines = []
        self.__hosts_map = {}

    def __add_host(self, host_line, force=True):
        if force or host_line.is_comment:
            self.__hosts_lines.append(host_line)

        if not host_line.is_comment:
            ip = host_line.ip
            domains = host_line.domains
            lines = self.__hosts_map.get(ip)
            if force:
                if not lines:
                    lines = []

                lines.append(host_line)
                self.__hosts_map[ip] = lines

                for domain in domains:
                    domain_lines = self.__hosts_map.get(domain, [])
                    domain_lines.append(host_line)
                    self.__hosts_map[domain] = domain_lines
            else:

                # self.__hosts_lines[]
                for domain in domains:
                    domain_lines = self.__hosts_map.get(domain, [])
                    # 说明域名被映射了多个ip
                    for domain_line in domain_lines:
                        if domain_line.ip == ip:
                            # 如果已经被映射，就删除当前host domain
                            host_line.remove_domain(domain)
                        else:
                            domain_line.remove_domain(domain)
                if host_line.is_comment:
                    print('%s 已存在' % host_line)
                    return

                self.__add_host(host_line, True)

    def parse(self, hosts_text):
        self.__hosts_lines = []
        self.__hosts_map.clear()

        hosts_lines = []
        for line in hosts_text.splitlines():
            hosts_lines.append(HostLine(line))

        for line in hosts_lines:
            self.__add_host(line)

    def add(self, line):
        self.__add_host(HostLine(line), False)

    def remove_domain(self, domain):
        for line in self.__hosts_map.get(domain, []):
            line.remove_domain(domain)

    def __str__(self):
        return '\n'.join(map(str, self.__hosts_lines))

    def to_hosts_text(self):
        return str(self)


def sys_hosts():
    system = platform.system()
    if system == 'Windows':
        return 'C:/Windows/System32/drivers/etc/hosts'
    elif system == 'Linux' or 'Darwin':
        return '/etc/hosts'


class HostsFile(Hosts):

    def __init__(self, path):
        super(HostsFile, self).__init__()
        self.__path = path

        with open(path) as fp:
            self.parse(fp.read())

    def save(self, path=None):
        path = self.__path if path is None else path
        with open(path, 'w') as fp:
            fp.write(str(self))


def sys_hostsfile():
    return HostsFile(sys_hosts())


def main():
    parser = argparse.ArgumentParser(description='hosts 修改脚本', usage='-h 查看帮助')

    parser.add_argument('-a', '--add-host', action='append',
                        help='添加host行, eg: -a "127.0.0.1 localhost"', default=[])

    parser.add_argument('-d', '--del-domains', action='append',
                        help='删除host, eg: -d "localhost"', default=[])

    parser.add_argument('-f', '--hosts-file',
                        help='hosts文件, eg, linux默认: -f "/etc/hosts"', default=sys_hosts())

    args = parser.parse_args()

    host_file = HostsFile(args.hosts_file)

    save = False

    if args.add_host:
        for line in args.add_host:
            host_file.add(line)
            save = True

    if args.del_domains:
        for domain in args.del_dominas:
            host_file.remove_domain(domain)
            save = True

    if save:
        host_file.save()

    print (host_file)


if __name__ == '__main__':
    main()
    print('执行完成!')
